### Beginning ####
rm(list = ls())

# Prediction plots
pacman::p_load(tidyverse, 
               ggeffects, 
               effects, 
               haven, 
               patchwork)

### Study 1 Prediction plots - unscaled variables ####

psych_long_us <- read_dta("./data_clean/psych_clean_long.dta") %>% 
  rename(StoryTrue = true,
         # CRT = zCRT_all,
         CRT = CRT_all,
         # AOT = AOT_all,
         AOT =  AOT_allUnscaled,
         # AOT_shortUnscaled,
         # AOT_longUnscaled,
         # RussianOrientation = russianOrientation,
         # ProRussiaPoliOrientation = zantiEuropeOrientation,
         RussianOrientation = XrussianOrientation,
         ProRussiaPoliOrientation = antiEuropeOrientation,
         Female = RSPSEX,
         Age = RSPAGE,
         Education = RSPEDUC,
         Income = income) %>% 
  mutate(AOT = (AOT - 1)/4)

load("./data_clean/main_clustered_models_study1_unscaled.Rdata")

### CRT RussianOrientation

quantile(psych_long_us$RussianOrientation)
quantile(psych_long_us$CRT)
range(psych_long_us$response)

mcrt1_study1_us <-
  ggpredict(main_models_us[[2]][[1]], 
            terms =  c("CRT [0.0000000, 0.1428571, 0.4285714, 0.5714286, 1.0000000]",
                       "RussianOrientation [-4, -1, 2]",
                       "StoryTrue [0, 1]"),
            vcov.fun = "vcovCL",
            vcov.type = "HC1",
            vcov.args =  dplyr::select(psych_long_us, Respondent_Serial, question)) %>%
  mutate(facet = ifelse(facet == 0, "Story False", "Story True"))

p1_study1_us <-
  ggplot(mcrt1_study1_us, aes(x, predicted, ymin = conf.low, ymax = conf.high, fill = group)) +
  geom_line(aes(color = group), size = 1) +
  geom_ribbon(alpha = 0.4) +
  geom_hline(aes(yintercept = 3.5), lty = "dashed") +
  facet_wrap(~facet) +
  scale_fill_discrete(name = "Russian\nOrientation", 
                      labels = c("1st Quartile (Low)", "Median",  "3rd Quartile (High)")
  ) +
  scale_color_discrete(guide = NULL) +
  scale_y_continuous(limits = c(1, 6),
                     breaks = seq(1, 6, 1)) +
  scale_x_continuous(breaks = seq(0, 1, 0.25),
                     labels = c("0", "0.25", "0.5", "0.75", "1")) +
  theme_bw() +
  theme(legend.position = "bottom") +
  labs(x = "Cognitive Reflection Test Score", y = "Predicted Belief");p1_study1_us

write.csv(mcrt1_study1_us %>% data.frame() %>% 
            rename(CRT = x, RussianOrientation = group), 
          file = "./tables/predictions_s1_CRT_rusori_.csv")

### AOT RussianOrientation

quantile(psych_long_us$AOT)

maot1_study1_us <-
  ggpredict(main_models_us[[2]][[2]], 
            terms =  c("AOT [0.0000000, 0.4062500, 0.4687500, 0.5208333, 1.000000]",
                       "RussianOrientation [-4, -1, 2]",
                       "StoryTrue [0, 1]"),
            vcov.fun = "vcovCL",
            vcov.type = "HC1",
            vcov.args =  dplyr::select(psych_long_us, Respondent_Serial, question)) %>%
  mutate(facet = ifelse(facet == 0, "Story False", "Story True"))

p2_study1_us <-
  ggplot(maot1_study1_us, aes(x, predicted, ymin = conf.low, ymax = conf.high, fill = group)) +
  geom_line(aes(color = group), size = 1) +
  geom_ribbon(alpha = 0.4) +
  geom_hline(aes(yintercept = 3.5), lty = "dashed") +
  facet_wrap(~facet) +
  scale_fill_discrete(name = "Russian\nOrientation", 
                      labels = c("1st Quartile (Low)", "Median",  "3rd Quartile (High)")
  ) +
  scale_y_continuous(limits = c(1, 6),
                     breaks = seq(1, 6, 1)) +
  scale_x_continuous(breaks = seq(0, 1, 0.25),
                     labels = c("0", "0.25", "0.5", "0.75", "1")) +
  scale_color_discrete(guide = NULL) +
  theme_bw() +
  theme(legend.position = "bottom") +
  labs(x = "Actively Open-Minded Test Score", y = "Predicted Belief");p2_study1_us

write.csv(maot1_study1_us %>% data.frame() %>% rename(CRT = x, RussianOrientation = group), file = "./tables/predictions_s1_AOT_rusori_.csv")

### CRT ProRussiaPoliOrientation

quantile(psych_long_us$ProRussiaPoliOrientation)
quantile(psych_long_us$CRT)

mcrt2_study1_us <-
  ggpredict(main_models_us[[2]][[3]], 
            terms =  c("CRT [0.0000000, 0.1428571, 0.4285714, 0.5714286, 1.0000000]",
                       "ProRussiaPoliOrientation [-1, 0, 1]", # label, pro-Russia, Anti-Russia, neither
                       "StoryTrue [0, 1]"),
            vcov.fun = "vcovCL",
            vcov.type = "HC1",
            vcov.args =  dplyr::select(psych_long_us, Respondent_Serial, question)) %>%
  mutate(facet = ifelse(facet == 0, "Story False", "Story True"))

p3_study1_us <-
  ggplot(mcrt2_study1_us, aes(x, predicted, ymin = conf.low, ymax = conf.high, fill = group)) +
  geom_line(aes(color = group), size = 1) +
  geom_ribbon(alpha = 0.4) +
  geom_hline(aes(yintercept = 3.5), lty = "dashed") +
  facet_wrap(~facet) +
  scale_fill_discrete(name = "Pro Russia\nPoli Orientation", 
                      labels = c("Anti-Russia", "Neutral", "Pro-Russia")
  ) +
  scale_y_continuous(limits = c(1, 6),
                     breaks = seq(1, 6, 1)) +
  scale_x_continuous(breaks = seq(0, 1, 0.25),
                     labels = c("0", "0.25", "0.5", "0.75", "1")) +
  scale_color_discrete(guide = NULL) +
  theme_bw() +
  theme(legend.position = "bottom") +
  labs(x = "Cognitive Reflection Test Score", y = "Predicted Belief");p3_study1_us

write.csv(mcrt2_study1_us %>% data.frame() %>% rename(CRT = x, ProRussiaPoliOrientation = group), file = "./tables/predictions_s1_CRT_poliori_.csv")

### AOT ProRussiaPoliOrientation

quantile(psych_long_us$AOT)

maot2_study1_us <-
  ggpredict(main_models_us[[2]][[4]], 
            terms =  c("AOT [0.000000, 0.4062500, 0.4687500, 0.5208333, 1.000000]",
                       "ProRussiaPoliOrientation [-1, 0, 1]",
                       "StoryTrue [0, 1]"),
            vcov.fun = "vcovCL",
            vcov.type = "HC1",
            vcov.args =  dplyr::select(psych_long_us, Respondent_Serial, question)) %>%
  mutate(facet = ifelse(facet == 0, "Story False", "Story True"))

p4_study1_us <-
  ggplot(maot2_study1_us, aes(x, predicted, ymin = conf.low, ymax = conf.high, fill = group)) +
  geom_line(aes(color = group), size = 1) +
  geom_ribbon(alpha = 0.4) +
  geom_hline(aes(yintercept = 3.5), lty = "dashed") +
  facet_wrap(~facet) +
  scale_fill_discrete(name = "Pro Russia\nPoli Orientation", 
                      labels = c("Anti-Russia", "Neutral", "Pro-Russia")
  ) +
  scale_y_continuous(limits = c(1, 6),
                     breaks = seq(1, 6, 1)) +
  scale_x_continuous(breaks = seq(0, 1, 0.25),
                     labels = c("0", "0.25", "0.5", "0.75", "1")) +
  scale_color_discrete(guide = NULL) +
  theme_bw() +
  theme(legend.position = "bottom") +
  labs(x = "Actively Open-Minded Test Score", y = "Predicted Belief");p4_study1_us

write.csv(maot2_study1_us %>% data.frame() %>% rename(CRT = x, ProRussiaPoliOrientation = group), file = "./tables/predictions_s1_AOT_poliori_.csv")

### S1 Grid Arrange

patch_s1_us <-
  (p1_study1_us + theme(legend.position = "right") + labs(x = NULL) +
     p2_study1_us + theme(legend.position = "none") + labs(y = NULL, x = NULL)) / 
  (p3_study1_us + theme(legend.position = "right") + 
     p4_study1_us + theme(legend.position = "none") + labs(y = NULL)); patch_s1_us

ggsave("./figs/s1_facet_us.png", patch_s1_us, width = 8.5, height = 9.5)

rm(psych_long_us)

### Study 2 Prediction Plots - unscaled variables ####

psych_long2_us <- read_dta("./data_clean/psych_clean_long2.dta") %>% 
  mutate(RSPEDUC = factor(RSPEDUC),
         INCOM = factor(INCOM),
         oblast = as_factor(oblast),
         narr_topic = as_factor(narr_topic)) %>% 
  rename(StoryTrue = narr_true,
         # CRT = zCRT_all,
         CRT = crt,
         # AOT = AOT_all,
         AOT = AOT_allUnscaled,
         # RussianOrientation = russianOrientation,
         RussianOrientation = russianOrientation_unscaled,
         # RussiaThreat = RussiaThreat_unscaled,
         RussiaNotAThreat = RussiaThreat,
         Female = RSPSEX,
         Age = RSPAGE,
         Education = RSPEDUC,
         Income = INCOM,
         OblastStrata = strtcode) %>% 
  mutate(CRT = CRT/4,
         AOT = (AOT - 1)/4)

load("./data_clean/main_clustered_models_study2_unscaled.Rdata")

### CRT RussianOrientation

quantile(psych_long2_us$CRT, na.rm = T)
quantile(psych_long2_us$RussianOrientation, na.rm = T)
range(psych_long2_us$narr_believe)

mcrt1_study2_us <- ggpredict(main_models_study2_us[[2]][[1]], 
                             terms =  c("CRT [0.00, 0.25, 0.5, 0.75, 1.00]",
                                        "RussianOrientation [-3, -1, 1]", 
                                        "StoryTrue [0, 1]"),
                   vcov.fun = "vcovCL",
                   vcov.type = "HC1",
                   vcov.args = dplyr::select(psych_long2_us, id, narrative)) %>%
  mutate(facet = ifelse(facet == 0, "Story False", "Story True"))

p1_study2_us <-
  ggplot(mcrt1_study2_us, aes(x, predicted + 4, ymin = conf.low + 4, ymax = conf.high + 4, fill = group)) +
  geom_line(aes(color = group), size = 1) +
  geom_ribbon(alpha = 0.4) +
  geom_hline(aes(yintercept = 4), lty = "dashed") +
  facet_wrap(~facet) +
  scale_fill_discrete(name = "Russian\nOrientation", 
                      labels = c("1st Quartile (Low)", "Median",  "3rd Quartile (High)")
  ) +
  scale_color_discrete(guide = NULL) +
  scale_y_continuous(limits = c(1, 7),
                     breaks = seq(1, 7, 1)) +
  scale_x_continuous(breaks = seq(0, 1, 0.25),
                     labels = c("0", "0.25", "0.5", "0.75", "1")) +
  theme_bw() +
  theme(legend.position = "bottom") +
  labs(x = "Cognitive Reflection Test Score", y = "Predicted Belief");p1_study2_us

write.csv(mcrt1_study2_us %>% data.frame(), file = "./tables/predictions_s2_CRT_rusori_.csv")

### AOT RussianOrientation

quantile(psych_long2_us$AOT, na.rm = T)

maot1_study2_us <-
  ggpredict(main_models_study2_us[[2]][[2]], 
            terms =  c("AOT [0.000000, 0.250000, 0.4166667, 0.5416667, 1.000000]",
                       "RussianOrientation [-3, -1, 1]", 
                       "StoryTrue [0, 1]"),
            vcov.fun = "vcovCL",
            vcov.type = "HC1",
            vcov.args = dplyr::select(psych_long2_us, id, narrative)) %>%
  mutate(facet = ifelse(facet == 0, "Story False", "Story True"))

p2_study2_us <-
  ggplot(maot1_study2_us, aes(x, predicted + 4, ymin = conf.low + 4, ymax = conf.high + 4, fill = group)) +
  geom_line(aes(color = group), size = 1) +
  geom_ribbon(alpha = 0.4) +
  geom_hline(aes(yintercept = 4), lty = "dashed") +
  facet_wrap(~facet) +
  scale_fill_discrete(name = "Russian\nOrientation", 
                      labels = c("1st Quartile (Low)", "Median",  "3rd Quartile (High)")
  ) +
  scale_y_continuous(limits = c(1, 7),
                     breaks = seq(1, 7, 1)) +
  scale_x_continuous(breaks = seq(0, 1, 0.25),
                     labels = c("0", "0.25", "0.5", "0.75", "1")) +
  scale_color_discrete(guide = NULL) +
  theme_bw() +
  theme(legend.position = "bottom") +
  labs(x = "Actively Open-Minded Test Score", y = "Predicted Belief");p2_study2_us

write.csv(maot1_study2_us %>% data.frame() %>% rename(CRT = x, RussianOrientation = group), file = "./tables/predictions_s2_AOT_rusori_.csv")


### CRT ProRussiaPoliOrientation

quantile(psych_long2_us$RussiaNotAThreat, na.rm = T)
quantile(psych_long2_us$CRT)
range(psych_long2_us$narr_believe)

mcrt2_study2_us <-
  ggpredict(main_models_study2_us[[2]][[3]], 
            terms =  c("CRT [0, 0.25, 0.5, 0.75, 1]",
                       "RussiaNotAThreat [-1, 0, 1]",
                       "StoryTrue [0, 1]"),
            vcov.fun = "vcovCL",
            vcov.type = "HC1",
            vcov.args = dplyr::select(psych_long2_us, id, narrative)) %>%
  mutate(facet = ifelse(facet == 0, "Story False", "Story True"))

p3_study2_us <-
  ggplot(mcrt2_study2_us, aes(x, predicted + 4, ymin = conf.low + 4, ymax = conf.high + 4, fill = group)) +
  geom_line(aes(color = group), size = 1) +
  geom_ribbon(alpha = 0.4) +
  geom_hline(aes(yintercept = 4), lty = "dashed") +
  facet_wrap(~facet) +
  scale_fill_discrete(name = "Russia Not\na Threat", 
                      labels = c("1st Quartile (High Threat)", "Median",  "3rd Quartile (Low Threat)")
  ) +
  scale_y_continuous(limits = c(1, 7),
                     breaks = seq(1, 7, 1)) +
  scale_x_continuous(breaks = seq(0, 1, 0.25),
                     labels = c("0", "0.25", "0.5", "0.75", "1")) +
  scale_color_discrete(guide = NULL) +
  theme_bw() +
  theme(legend.position = "bottom") +
  labs(x = "Cognitive Reflection Test Score", y = "Predicted Belief");p3_study2_us

write.csv(mcrt2_study2_us %>% data.frame() %>% rename(CRT = x, ProRussiaPoliOrientation = group), file = "./tables/predictions_s2_CRT_poliori_.csv")

### AOT ProRussiaPoliOrientation

quantile(psych_long2_us$AOT, na.rm = T)

maot2_study2_us <-
  ggpredict(main_models_study2_us[[2]][[4]], 
            terms =  c("AOT [0.000000, 0.250000, 0.4166667, 0.5416667, 1.000000]",
                       "RussiaNotAThreat [-1, 0, 1]",
                       "StoryTrue [0, 1]"),
            vcov.fun = "vcovCL",
            vcov.type = "HC1",
            vcov.args = dplyr::select(psych_long2_us, id, narrative)) %>%
  mutate(facet = ifelse(facet == 0, "Story False", "Story True"))

p4_study2_us <-
  ggplot(maot2_study2_us, aes(x, predicted + 4, ymin = conf.low + 4, ymax = conf.high + 4, fill = group)) +
  geom_line(aes(color = group), size = 1) +
  geom_ribbon(alpha = 0.4) +
  geom_hline(aes(yintercept = 4), lty = "dashed") +
  facet_wrap(~facet) +
  scale_fill_discrete(name = "Russia Not\na Threat", 
                      labels = c("1st Quartile (High Threat)", "Median",  "3rd Quartile (Low Threat)")
  ) +
  scale_y_continuous(limits = c(1, 7),
                     breaks = seq(1, 7, 1)) +
  scale_x_continuous(breaks = seq(0, 1, 0.25),
                     labels = c("0", "0.25", "0.5", "0.75", "1")) +
  scale_color_discrete(guide = NULL) +
  theme_bw() +
  theme(legend.position = "bottom") +
  labs(x = "Actively Open-Minded Test Score", y = "Predicted Belief");p4_study2_us

write.csv(maot2_study2_us %>% data.frame() %>% rename(CRT = x, ProRussiaPoliOrientation = group), file = "./tables/predictions_s2_AOT_poliori_.csv")

### S2 Grid Arrange

patch_s2_us <-
  (p1_study2_us + theme(legend.position = "right") + labs(x = NULL) +
     p2_study2_us + theme(legend.position = "none") + labs(y = NULL, x = NULL)) / 
  (p3_study2_us + theme(legend.position = "right") + 
     p4_study2_us + theme(legend.position = "none") + labs(y = NULL)) + plot_layout();patch_s2_us

ggsave("./figs/s2_facet_us.png", patch_s2_us, width = 8.5, height = 9.5)

rm(psych_long2_us)
